"""
统一加密解密工具
================

为所有微服务提供统一的数据加密、解密、签名验证和密钥管理功能。
支持对称加密、非对称加密、数字签名和安全的密钥存储。

特性：
- AES对称加密
- RSA非对称加密
- 数字签名和验证
- 安全的密钥生成和管理
- 数据完整性校验
- 时间戳签名
"""

from datetime import datetime, timedelta
from typing import Dict, Any, Optional, Tuple, Union
import json
import logging

from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
import hashlib
import hmac
import secrets

from ...common.config import get_service_config
logger = logging.getLogger(__name__)


class UnifiedCryptoManager:
    """统一加密管理器"""
    
    def __init__(self, service_name: str = "default"):
        """
        初始化加密管理器
        
        Args:
            service_name: 服务名称，用于获取特定配置
        """
        self.service_name = service_name
        self.config = get_service_config(service_name)
        self.crypto_config = self.config.get("crypto", {})
        
        # 获取主密钥（从环境变量或配置）
        self.master_key = self._get_master_key()
        
        # 初始化Fernet实例
        self.fernet = Fernet(self.master_key)
        
        # RSA密钥对缓存
        self._rsa_private_key = None
        self._rsa_public_key = None
    
    def _get_master_key(self) -> bytes:
        """获取主加密密钥"""
        import os
        
        # 优先从环境变量获取
        key_env = f"{self.service_name.upper()}_MASTER_KEY"
        master_key = os.getenv(key_env)
        
        if master_key:
            # 如果是base64编码的密钥
            try:
                return base64.urlsafe_b64decode(master_key)
            except Exception:
                # 如果不是base64，则从字符串生成密钥
                return self._derive_key_from_password(master_key)
        
        # 从配置文件获取
        config_key = self.crypto_config.get("master_key")
        if config_key:
            try:
                return base64.urlsafe_b64decode(config_key)
            except Exception:
                return self._derive_key_from_password(config_key)
        
        # 生成默认密钥（仅用于开发环境）
        logger.warning("Using default master key for development. DO NOT use in production!")
        return Fernet.generate_key()
    
    def _derive_key_from_password(self, password: str, salt: Optional[bytes] = None) -> bytes:
        """从密码派生密钥"""
        if salt is None:
            salt = b"redfire_default_salt"  # 在生产环境中应该使用随机盐
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        
        return base64.urlsafe_b64encode(kdf.derive(password.encode()))
    
    def encrypt_data(self, data: Union[str, bytes, dict]) -> str:
        """
        加密数据
        
        Args:
            data: 要加密的数据
            
        Returns:
            str: base64编码的加密数据
        """
        try:
            # 将数据转换为字节
            if isinstance(data, dict):
                data_bytes = json.dumps(data).encode('utf-8')
            elif isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 使用Fernet加密
            encrypted = self.fernet.encrypt(data_bytes)
            
            # 返回base64编码的结果
            return base64.urlsafe_b64encode(encrypted).decode('utf-8')
            
        except Exception as e:
            logger.error(f"Data encryption failed: {e}")
            raise
    
    def decrypt_data(self, encrypted_data: str, return_type: str = "str") -> Union[str, bytes, dict]:
        """
        解密数据
        
        Args:
            encrypted_data: base64编码的加密数据
            return_type: 返回类型 ("str", "bytes", "dict")
            
        Returns:
            Union[str, bytes, dict]: 解密后的数据
        """
        try:
            # 解码base64
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_data.encode('utf-8'))
            
            # 使用Fernet解密
            decrypted_bytes = self.fernet.decrypt(encrypted_bytes)
            
            # 根据返回类型转换
            if return_type == "bytes":
                return decrypted_bytes
            elif return_type == "dict":
                return json.loads(decrypted_bytes.decode('utf-8'))
            else:
                return decrypted_bytes.decode('utf-8')
                
        except Exception as e:
            logger.error(f"Data decryption failed: {e}")
            raise
    
    def encrypt_with_timestamp(self, data: Union[str, dict], ttl_seconds: Optional[int] = None) -> str:
        """
        带时间戳的加密
        
        Args:
            data: 要加密的数据
            ttl_seconds: 生存时间（秒）
            
        Returns:
            str: 加密后的数据
        """
        try:
            # 添加时间戳信息
            timestamp_data = {
                "data": data,
                "timestamp": datetime.utcnow().isoformat(),
                "ttl": ttl_seconds
            }
            
            return self.encrypt_data(timestamp_data)
            
        except Exception as e:
            logger.error(f"Timestamp encryption failed: {e}")
            raise
    
    def decrypt_with_timestamp(self, encrypted_data: str) -> Tuple[Any, bool]:
        """
        带时间戳的解密
        
        Args:
            encrypted_data: 加密的数据
            
        Returns:
            Tuple[Any, bool]: (解密的数据, 是否在有效期内)
        """
        try:
            # 解密数据
            timestamp_data = self.decrypt_data(encrypted_data, return_type="dict")
            
            # 检查时间戳
            timestamp = datetime.fromisoformat(timestamp_data["timestamp"])
            ttl = timestamp_data.get("ttl")
            
            is_valid = True
            if ttl:
                expiry = timestamp + timedelta(seconds=ttl)
                is_valid = datetime.utcnow() <= expiry
            
            return timestamp_data["data"], is_valid
            
        except Exception as e:
            logger.error(f"Timestamp decryption failed: {e}")
            raise
    
    def generate_rsa_keys(self, key_size: int = 2048) -> Tuple[bytes, bytes]:
        """
        生成RSA密钥对
        
        Args:
            key_size: 密钥长度
            
        Returns:
            Tuple[bytes, bytes]: (私钥, 公钥)
        """
        try:
            # 生成私钥
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=key_size,
                backend=default_backend()
            )
            
            # 获取公钥
            public_key = private_key.public_key()
            
            # 序列化私钥
            private_pem = private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
            
            # 序列化公钥
            public_pem = public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )
            
            self._rsa_private_key = private_key
            self._rsa_public_key = public_key
            
            return private_pem, public_pem
            
        except Exception as e:
            logger.error(f"RSA key generation failed: {e}")
            raise
    
    def load_rsa_keys(self, private_key_pem: bytes, public_key_pem: bytes):
        """
        加载RSA密钥对
        
        Args:
            private_key_pem: PEM格式的私钥
            public_key_pem: PEM格式的公钥
        """
        try:
            self._rsa_private_key = serialization.load_pem_private_key(
                private_key_pem,
                password=None,
                backend=default_backend()
            )
            
            self._rsa_public_key = serialization.load_pem_public_key(
                public_key_pem,
                backend=default_backend()
            )
            
        except Exception as e:
            logger.error(f"RSA key loading failed: {e}")
            raise
    
    def rsa_encrypt(self, data: Union[str, bytes], use_public_key: bool = True) -> str:
        """
        RSA加密
        
        Args:
            data: 要加密的数据
            use_public_key: 是否使用公钥加密
            
        Returns:
            str: base64编码的加密数据
        """
        try:
            if not self._rsa_public_key or not self._rsa_private_key:
                raise ValueError("RSA keys not loaded")
            
            # 转换数据
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 选择密钥
            key = self._rsa_public_key if use_public_key else self._rsa_private_key
            
            # RSA加密
            encrypted = key.encrypt(
                data_bytes,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            
            return base64.urlsafe_b64encode(encrypted).decode('utf-8')
            
        except Exception as e:
            logger.error(f"RSA encryption failed: {e}")
            raise
    
    def rsa_decrypt(self, encrypted_data: str, use_private_key: bool = True) -> str:
        """
        RSA解密
        
        Args:
            encrypted_data: base64编码的加密数据
            use_private_key: 是否使用私钥解密
            
        Returns:
            str: 解密后的数据
        """
        try:
            if not self._rsa_public_key or not self._rsa_private_key:
                raise ValueError("RSA keys not loaded")
            
            # 解码base64
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_data.encode('utf-8'))
            
            # 选择密钥
            key = self._rsa_private_key if use_private_key else self._rsa_public_key
            
            # RSA解密
            decrypted = key.decrypt(
                encrypted_bytes,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            
            return decrypted.decode('utf-8')
            
        except Exception as e:
            logger.error(f"RSA decryption failed: {e}")
            raise
    
    def sign_data(self, data: Union[str, bytes]) -> str:
        """
        数字签名
        
        Args:
            data: 要签名的数据
            
        Returns:
            str: base64编码的签名
        """
        try:
            if not self._rsa_private_key:
                raise ValueError("RSA private key not loaded")
            
            # 转换数据
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 生成签名
            signature = self._rsa_private_key.sign(
                data_bytes,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            
            return base64.urlsafe_b64encode(signature).decode('utf-8')
            
        except Exception as e:
            logger.error(f"Data signing failed: {e}")
            raise
    
    def verify_signature(self, data: Union[str, bytes], signature: str) -> bool:
        """
        验证数字签名
        
        Args:
            data: 原始数据
            signature: base64编码的签名
            
        Returns:
            bool: 签名是否有效
        """
        try:
            if not self._rsa_public_key:
                raise ValueError("RSA public key not loaded")
            
            # 转换数据
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 解码签名
            signature_bytes = base64.urlsafe_b64decode(signature.encode('utf-8'))
            
            # 验证签名
            self._rsa_public_key.verify(
                signature_bytes,
                data_bytes,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            
            return True
            
        except Exception as e:
            logger.warning(f"Signature verification failed: {e}")
            return False
    
    def generate_hmac(self, data: Union[str, bytes], key: Optional[bytes] = None) -> str:
        """
        生成HMAC
        
        Args:
            data: 要签名的数据
            key: HMAC密钥（可选，默认使用主密钥）
            
        Returns:
            str: hex编码的HMAC
        """
        try:
            # 转换数据
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 使用密钥
            if key is None:
                key = self.master_key
            
            # 生成HMAC
            mac = hmac.new(key, data_bytes, hashlib.sha256)
            return mac.hexdigest()
            
        except Exception as e:
            logger.error(f"HMAC generation failed: {e}")
            raise
    
    def verify_hmac(self, data: Union[str, bytes], expected_hmac: str, key: Optional[bytes] = None) -> bool:
        """
        验证HMAC
        
        Args:
            data: 原始数据
            expected_hmac: 期望的HMAC值
            key: HMAC密钥（可选）
            
        Returns:
            bool: HMAC是否有效
        """
        try:
            computed_hmac = self.generate_hmac(data, key)
            return hmac.compare_digest(computed_hmac, expected_hmac)
        except Exception as e:
            logger.error(f"HMAC verification failed: {e}")
            return False
    
    def generate_secure_token(self, length: int = 32) -> str:
        """
        生成安全令牌
        
        Args:
            length: 令牌长度
            
        Returns:
            str: URL安全的base64编码令牌
        """
        return secrets.token_urlsafe(length)
    
    def hash_data(self, data: Union[str, bytes], algorithm: str = "sha256") -> str:
        """
        哈希数据
        
        Args:
            data: 要哈希的数据
            algorithm: 哈希算法
            
        Returns:
            str: hex编码的哈希值
        """
        try:
            # 转换数据
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            else:
                data_bytes = data
            
            # 选择算法
            if algorithm == "sha256":
                hasher = hashlib.sha256()
            elif algorithm == "sha512":
                hasher = hashlib.sha512()
            elif algorithm == "md5":
                hasher = hashlib.md5()
            else:
                raise ValueError(f"Unsupported hash algorithm: {algorithm}")
            
            hasher.update(data_bytes)
            return hasher.hexdigest()
            
        except Exception as e:
            logger.error(f"Data hashing failed: {e}")
            raise
    
    def constant_time_compare(self, a: str, b: str) -> bool:
        """
        常量时间比较（防止时序攻击）
        
        Args:
            a: 字符串A
            b: 字符串B
            
        Returns:
            bool: 是否相等
        """
        return hmac.compare_digest(a, b)


# =============================================================================
# 全局加密管理器实例
# =============================================================================

_crypto_managers = {}


def get_crypto_manager(service_name: str = "default") -> UnifiedCryptoManager:
    """
    获取加密管理器实例
    
    Args:
        service_name: 服务名称
        
    Returns:
        UnifiedCryptoManager: 加密管理器实例
    """
    if service_name not in _crypto_managers:
        _crypto_managers[service_name] = UnifiedCryptoManager(service_name)
    return _crypto_managers[service_name]


# =============================================================================
# 便捷函数
# =============================================================================

def encrypt_data(data: Union[str, bytes, dict], service_name: str = "default") -> str:
    """
    加密数据（便捷函数）
    
    Args:
        data: 要加密的数据
        service_name: 服务名称
        
    Returns:
        str: 加密后的数据
    """
    manager = get_crypto_manager(service_name)
    return manager.encrypt_data(data)


def decrypt_data(encrypted_data: str, return_type: str = "str", service_name: str = "default") -> Union[str, bytes, dict]:
    """
    解密数据（便捷函数）
    
    Args:
        encrypted_data: 加密的数据
        return_type: 返回类型
        service_name: 服务名称
        
    Returns:
        Union[str, bytes, dict]: 解密后的数据
    """
    manager = get_crypto_manager(service_name)
    return manager.decrypt_data(encrypted_data, return_type)


def generate_secure_token(length: int = 32, service_name: str = "default") -> str:
    """
    生成安全令牌（便捷函数）
    
    Args:
        length: 令牌长度
        service_name: 服务名称
        
    Returns:
        str: 安全令牌
    """
    manager = get_crypto_manager(service_name)
    return manager.generate_secure_token(length)


def hash_data(data: Union[str, bytes], algorithm: str = "sha256", service_name: str = "default") -> str:
    """
    哈希数据（便捷函数）
    
    Args:
        data: 要哈希的数据
        algorithm: 哈希算法
        service_name: 服务名称
        
    Returns:
        str: 哈希值
    """
    manager = get_crypto_manager(service_name)
    return manager.hash_data(data, algorithm)


def generate_hmac(data: Union[str, bytes], key: Optional[bytes] = None, service_name: str = "default") -> str:
    """
    生成HMAC（便捷函数）
    
    Args:
        data: 要签名的数据
        key: HMAC密钥
        service_name: 服务名称
        
    Returns:
        str: HMAC值
    """
    manager = get_crypto_manager(service_name)
    return manager.generate_hmac(data, key)


def verify_hmac(data: Union[str, bytes], expected_hmac: str, key: Optional[bytes] = None, service_name: str = "default") -> bool:
    """
    验证HMAC（便捷函数）
    
    Args:
        data: 原始数据
        expected_hmac: 期望的HMAC值
        key: HMAC密钥
        service_name: 服务名称
        
    Returns:
        bool: HMAC是否有效
    """
    manager = get_crypto_manager(service_name)
    return manager.verify_hmac(data, expected_hmac, key)


# =============================================================================
# 导出的主要函数
# =============================================================================

__all__ = [
    "UnifiedCryptoManager",
    "get_crypto_manager",
    "encrypt_data",
    "decrypt_data",
    "generate_secure_token",
    "hash_data",
    "generate_hmac",
    "verify_hmac"
]
